module class_solver

    use class_input, only: indat,uout
    use class_dat, only: rhsdat,coefdat,flux,keff
    use class_2d_matx

    implicit none

    contains

    subroutine solver(inp,A,phi,rhs)

        implicit none
#include "finclude/petsc.h90"

        ! input arguments
        type(indat), intent(in)                  :: inp
        type(rhsdat), intent(in)                 :: rhs
        Mat, dimension(:), pointer               :: A
        type(flux), intent(out)                  :: phi
        ! local
        KSP               :: ksp
        PC                :: Prec
        PetscInt          :: it, nrow, mxit, g, gg, ierr, i
        PetscScalar       :: phieps, phierr, tol, zero, one, none, norm, tmpnorm, &
                             errk, epsk, errs, epss, keff0, gain, loss
        Vec               :: scsrc, fisrc, tmp, fisrc0
        Vec, dimension(:), pointer  :: phi0
        logical           :: ldbug

        ldbug = .false.
        zero  = 0.0
        one   = 1.0
        none  = -1.0

        nrow = inp%nxfm*inp%nyfm

        ! get the flux vectors ready
        allocate( phi%val(inp%numg), phi0(inp%numg) )
        call VecDuplicate( rhs%nuSigF(1), scsrc , ierr )
        call VecDuplicate( rhs%nuSigF(1), tmp , ierr )
        call VecDuplicate( rhs%nuSigF(1), fisrc , ierr )
        call VecDuplicate( rhs%nuSigF(1), fisrc0 , ierr )
        do g = 1, inp%numg
            call VecDuplicate( rhs%nuSigF(1), phi%val(g) , ierr )
            call VecDuplicate( rhs%nuSigF(1), phi0(g) , ierr )
            call VecSet( phi%val(g), zero, ierr )
            call VecSet( phi0(g), zero, ierr )
        end do

        !-------------------------------------------------
        ! initialize the solver context
        !-------------------------------------------------

        ! create solver context
        call KSPCreate( PETSC_COMM_WORLD, ksp, ierr )
        tol = 1.d-7
        call KSPSetTolerances(ksp,tol,PETSC_DEFAULT_DOUBLE_PRECISION,     &
                              PETSC_DEFAULT_DOUBLE_PRECISION,             &
                              PETSC_DEFAULT_INTEGER,ierr)
        call KSPSetType(ksp,KSPGMRES,ierr) ! using gmres
        call KSPSetFromOptions(ksp,ierr)   ! enable command line parameters

        ! create preconditioner
        call KSPGetPC( ksp, Prec, ierr )
        call PCSetType( Prec, PCILU, ierr ) ! essentially override gmres w/ g.elim.

        ! solver parameters should be in user input
        phieps = 1e-10
        mxit   = 20
        it     = 0
        phierr = 1
        norm   = 1.0

       !-------------------------------------------------------
       ! FIXED SOURCE WITHOUT MULTIPLICATION
       !-------------------------------------------------------
       if ( inp%ptype .eq. 0 ) then

        do while (norm > phieps .and. it < mxit)

            ! keep old solution for comparision
            do g = 1, inp%numg
                call VecCopy( phi%val(g), phi0(g), ierr )
            end do

            do g = 1, inp%numg

                if ( ldbug ) write (*,*) "group ",g
                if ( ldbug ) call MatView(A(g),PETSC_VIEWER_STDOUT_SELF,ierr)

                ! set group operator
                call KSPSetOperators( ksp, A(g), A(g), SAME_NONZERO_PATTERN, ierr )

                if (g.eq.1) then

                    call KSPSolve( ksp, rhs%fixedSource(g), phi%val(g), ierr )

                else

                    ! reset scattering source
                    call VecSet(  scsrc, zero, ierr )
                    call VecSet(  tmp, zero, ierr )
                    if ( ldbug ) write (*,*) " *** scatter vector initial: "
                    if ( ldbug ) call VecView(scsrc,PETSC_VIEWER_STDOUT_SELF,ierr)
                    do gg = 1, inp%numg ! go over all groups; could make upscatter an option
                        if ( ldbug ) write (*,*) " gg = ", gg, " g = ",g
                        if ( ldbug ) write (*,*) " *** tmp before scatter vector: "
                        if ( ldbug ) call VecView(rhs%sigScatter(g,gg),PETSC_VIEWER_STDOUT_SELF,ierr)
                        ! VecPointwiseMult(w,x,y) --> w = x.*y
                        call VecPointwiseMult( tmp, phi%val(gg), rhs%sigScatter(g,gg), ierr )
                        call VecAXPY( scsrc, one, tmp, ierr )
                    end do
                    if ( ldbug ) write (*,*) " *** scatter vector: "
                    call VecView(scsrc,PETSC_VIEWER_STDOUT_SELF,ierr)
                    ! need to build temporary rhs vector
                    call VecWAXPY( tmp, one, rhs%fixedSource(g), scsrc, ierr )
                    if ( ldbug ) write (*,*) " *** temp rhs vector: "
                    if ( ldbug ) call VecView(tmp,PETSC_VIEWER_STDOUT_SELF,ierr)

                    ! solve the system
                    call KSPSolve( ksp, tmp, phi%val(g), ierr )

                end if

				! i think I need to release the current A(g) from ksp
                !call KSPSetOperators( ksp, A(g), A(g), SAME_NONZERO_PATTERN, ierr )

            end do

            ! make a function returning norm given phi0 and phi
            norm = 0.0
            do g = 1, inp%numg
                call VecWAXPY( tmp, none, phi0(g), phi%val(g), ierr )
                call VecNorm( tmp, NORM_2, tmpnorm, ierr )
                norm = tmpnorm**2
            end do
            norm = sqrt(norm)
            write(*,*) "iteration: ",it, " norm: ",norm
            it = it + 1;

        end do

       !-------------------------------------------------------
       ! FIXED SOURCE WITH MULTIPLICATION -- REQUIRES keff
       !-------------------------------------------------------
       elseif ( inp%ptype .eq. -1 ) then
        write (*,*) " fixed source with mult "
        call VecSet(  fisrc, zero, ierr )
        keff = 1.0 ! this needs to be variable or user-defined eventually

        do while (norm > phieps .and. it < mxit)

            ! keep old solution for comparision
            do g = 1, inp%numg
                call VecCopy( phi%val(g), phi0(g), ierr )
            end do

            do g = 1, inp%numg ! groups

                if ( ldbug ) write (*,*) "group ",g
                if ( ldbug ) call MatView(A(g),PETSC_VIEWER_STDOUT_SELF,ierr)

                ! set group operator
                call KSPSetOperators( ksp, A(g), A(g), SAME_NONZERO_PATTERN, ierr )

                call VecSet(  tmp, zero, ierr )

                if (g.eq.1) then

                    ! phi(:,1) = A{1} \ ( Si(:,1) + chi(:,1).*s(:,1)/keff ) ;
                    ! creat tmp to be Si + chi.*s/keff
                    call VecPointwiseMult( tmp, rhs%xiSpect(g), fisrc, ierr )
                    write (*,*) " tmp 1 ! = "
                    call VecView(tmp,PETSC_VIEWER_STDOUT_SELF,ierr)
                    ! VecAYPX(y,a,x) -> y = a*y+x
                    call VecAYPX( tmp, one/keff, rhs%fixedSource(g), ierr )
                    write (*,*) " tmp 2  = "
                    call VecView(tmp,PETSC_VIEWER_STDOUT_SELF,ierr)
                    write (*,*) " fixed   = "
                    call VecView(rhs%fixedSource(g),PETSC_VIEWER_STDOUT_SELF,ierr)
                    call KSPSolve( ksp, tmp, phi%val(g), ierr)
                    call VecView( phi%val(g),PETSC_VIEWER_STDOUT_SELF,ierr)

                else

                    ! reset scattering source
                    call VecSet(  scsrc, zero, ierr )
                    do gg = 1, inp%numg ! go over all groups; could make upscatter an option
                        call VecPointwiseMult( tmp, phi%val(gg), rhs%sigScatter(g,gg), ierr )
                        call VecAXPY( scsrc, one, tmp, ierr )
                    end do
                    ! make xi*fissrc
                    call VecPointwiseMult( tmp, rhs%xiSpect(g), fisrc, ierr )
                    write (*,*) " xi*fissrc "
                    call VecView( tmp,PETSC_VIEWER_STDOUT_SELF,ierr)
                    ! make xi*fissrc/k + fixedsrc
                    call VecAYPX( tmp, one/keff, rhs%fixedSource(g), ierr )
                    write (*,*) " xi*fissrc/k + fixedsrc "
                    call VecView( tmp,PETSC_VIEWER_STDOUT_SELF,ierr)
                    ! make xi*fissrc/k + fixedsrc + sctsrc
                    call VecAYPX( tmp, one, scsrc, ierr )
                    ! solve the system
                    call KSPSolve( ksp, tmp, phi%val(g), ierr )
                    write (*,*) " xi*fissrc/k + fixedsrc + sctsrc "
                    call VecView( tmp,PETSC_VIEWER_STDOUT_SELF,ierr)
                    write (*,*) " phi(2) "
                    call VecView( phi%val(g),PETSC_VIEWER_STDOUT_SELF,ierr)

                end if

				! i think I need to release the current A(g) from ksp
                !call KSPSetOperators( ksp, A(g), A(g), SAME_NONZERO_PATTERN, ierr )

            end do ! groups

            ! reset fission source
            !s(:,1) = 0; % reset fission source
            call VecSet(  fisrc, zero, ierr )
            !for g = 1:numg
            do g = 1, inp%numg
            !   s(:,1) = s(:,1) + Sh(:,g).*phi(:,g);
                call VecPointwiseMult( tmp, rhs%nuSigF(g), phi%val(g), ierr )
                call VecAYPX( fisrc, one, tmp, ierr )
            !end
            end do

            ! make a function returning norm given phi0 and phi
            norm = 0.0
            do g = 1, inp%numg
                call VecWAXPY( tmp, none, phi0(g), phi%val(g), ierr )
                call VecNorm( tmp, NORM_2, tmpnorm, ierr )
                norm = tmpnorm**2
            end do
            norm = sqrt(norm)
            write(*,*) "iteration: ",it, " norm: ",norm
            it = it + 1;

        end do ! while

       !-------------------------------------------------------
       ! EIGENVALUE
       !-------------------------------------------------------
       else
        write (uout,*) " beginning eigenvalue iterations "
        call VecSet(  fisrc, one, ierr ) ! uniform initial guess
        call VecSet(  fisrc0, zero, ierr ) ! to store last source

        keff    = 1.0 ! perhaps user gives
        keff0   = 0.0
        errk    = 1.0
        errs    = 1.0
        epsk    = inp%epsk
        epss    = inp%epss
        it      = 0
        mxit   = inp%maxit

		!    while  ( ( errK>epsK ) || (errS>epsS) ) && iter < itmx
        do while ( (errk .gt. epsk .or. errs .gt. epss) .and. it .lt. mxit )

            do g = 1, inp%numg ! groups

                ! set group operator
                call KSPSetOperators( ksp, A(g), A(g), SAME_NONZERO_PATTERN, ierr )

                call VecSet(  tmp, zero, ierr )

                if (g.eq.1) then

                    call VecPointwiseMult( tmp, rhs%xiSpect(g), fisrc, ierr )
                    call VecScale( tmp, one/keff, ierr )
                    call KSPSolve( ksp, tmp, phi%val(g), ierr)

                else

                    ! reset scattering source
                    call VecSet(  scsrc, zero, ierr )
                    do gg = 1, inp%numg ! go over all groups; could make upscatter an option
                        call VecPointwiseMult( tmp, phi%val(gg), rhs%sigScatter(g,gg), ierr )
                        call VecAXPY( scsrc, one, tmp, ierr )
                    end do
                    ! make xi*fissrc
                    call VecPointwiseMult( tmp, rhs%xiSpect(g), fisrc, ierr )
                    ! make xi*fissrc/k
                    call VecScale( tmp, one/keff, ierr )
                    ! make xi*fissrc/k + sctsrc
                    call VecAYPX( tmp, one, scsrc, ierr )
                    ! solve the system
                    call KSPSolve( ksp, tmp, phi%val(g), ierr )

                end if

            end do ! groups

            ! save fission source and keff
            call VecCopy( fisrc, fisrc0, ierr )
            keff0 = keff
            ! reset fission source
            call VecSet(  fisrc, zero, ierr )
            do g = 1, inp%numg
                call VecPointwiseMult( tmp, rhs%nuSigF(g), phi%val(g), ierr )
                call VecAYPX( fisrc, one, tmp, ierr )
            end do
            ! compute new keff
            call VecDot( inp%dv, fisrc, gain, ierr )
            call VecDot( inp%dv, fisrc0, loss, ierr )
            keff = keff0*gain/loss
            ! find max fisrc error
            call VecWAXPY( tmp, none, fisrc, fisrc0, ierr )
            call VecNorm( tmp, NORM_1, errs, ierr )
            errk = abs( (keff-keff0)/keff )

            ! print iteration information
            if (inp%plotout.ne.0) write(uout,*) "iteration: ",it, " errs: ",errs, " errk: ",errk, " keff: ",keff

            it = it + 1;

        end do ! while

       end if

        ! debug flux prints
        if ( ldbug ) write (*,*) " phi group 1 = "
        if ( ldbug ) call VecView(phi%val(1),PETSC_VIEWER_STDOUT_WORLD,ierr)
        if ( ldbug ) write (*,*) " phi group 2 = "
        if ( ldbug ) call VecView(phi%val(2),PETSC_VIEWER_STDOUT_WORLD,ierr)
        if ( ldbug ) write (*,*) " *** final keff = ",keff," *** "
        ! destroy temporary vectors
        do g = 1, inp%numg
            call VecDestroy(phi0(g),ierr)
        end do
        call VecDestroy(scsrc,ierr)
        call VecDestroy(tmp,ierr)
        call VecDestroy(fisrc,ierr)


    end subroutine solver

end module class_solver
